from paddle.io import Dataset,DataLoader
import xml.etree.ElementTree as ET
import glob
import cv2
import numpy as np
classesName = ["holothurian", "echinus", "scallop", "starfish"]


class myDataSet(Dataset):
    def __init__(self, url, anchors):
        super(myDataSet, self).__init__()
        self.imgUrl = glob.glob(url+'\\image\\*.jpg')
        self.boxUrl = glob.glob(url+'\\box\\*.xml')
        self.anchors = anchors

    def __getitem__(self, index):
        image = cv2.imread(self.imgUrl[index], -1)
        box = getbox(self.boxUrl[index])
        image, box = transImg(image, box)
        box = true_boxes(box, self.anchors)  # (3,size,size,3,5+num)
        image = image.transpose(2, 0, 1)  # (3,416,416)
        image = image/255
        return image, box[0], box[1], box[2]

    def __len__(self):
        return len(self.imgUrl)

# 图像缩放
def transImg(img,box,outshape=(416, 416)):
    w, h, _ = img.shape
    scale = outshape[0]/w if outshape[0]/w < outshape[1]/h else outshape[1]/h
    img = cv2.resize(img, fx=scale, fy=scale, dsize=(0, 0))
    if outshape[0]/w == scale:
        p = outshape[1] - img.shape[1]
        img = np.pad(img, ((0, 0), (p//2, p-p//2), (0, 0)), 'constant')
        for i in range(len(box)):
            box[i][0] = box[i][0]*scale+p//2
            box[i][1] = box[i][1]*scale
            box[i][2] = box[i][2]*scale+p//2
            box[i][3] = box[i][3]*scale
    else:
        p = outshape[0] - img.shape[0]
        img = np.pad(img, ((p//2, p-p//2), (0, 0), (0, 0)), 'constant')
        for i in range(len(box)):
            box[i][0] = box[i][0]*scale
            box[i][1] = box[i][1]*scale+p//2
            box[i][2] = box[i][2]*scale
            box[i][3] = box[i][3]*scale+p//2
    return img, box

# 得到box
def getbox(url):
    tree = ET.parse(url)
    root = tree.getroot()
    box = []
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classesName:
            continue
        cls_id = classesName.index(cls)
        xmlbox = obj.find('bndbox')
        b = [int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text),
             int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text)]
        b.append(cls_id)
        box.append(b)
    return box

# 得到label  box的值被归一化到  0-1
def true_boxes(box,anchors,shape=(416,416),num_classes=4):
    box = np.array(box)  # (n,5)

    grid_shapes = [shape[0] // [32, 16, 8][l] for l in range(3)]  # (13,26,52)
    y_true = [np.zeros((grid_shapes[l], grid_shapes[l], 3, 5+num_classes),
        dtype='float32') for l in range(3)]    # (3,size,size,3,5+num)
    if len(box) == 0:
        return y_true
    anchor_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
    box_xy = (box[:, 0:2]+box[:, 2:4])//2
    box_wh = box[:, 2:4] - box[:, 0:2]
    box[:, 0:2] = box_xy/shape[0]  # 坐标轴归一化
    box[:, 2:4] = box_wh/shape[0]

    anchors = np.array(anchors)  # (9,2)
    anchors = np.expand_dims(anchors, 0)  # (1,9,2)
    # 为了求iou,将只有wh的anchor放在中心点为(0,0)的位置，
    anchor_maxes = anchors / 2.  # 右下坐标
    anchor_mins = -anchor_maxes  # 左上坐标 (1,9,2)
    wh = np.expand_dims(box_wh, 1)  # (n,1,2)
    box_maxes = wh/2
    box_mins = -box_maxes
    intersect_mins = np.maximum(box_mins, anchor_mins)  # (n,9,2)
    intersect_maxes = np.minimum(box_maxes, anchor_maxes)  # (n,9,2)
    intersect_wh = np.maximum(intersect_maxes - intersect_mins, 0.)  # (n,9,2)
    intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1]  # (n,9)
    box_area = wh[..., 0] * wh[..., 1]  # (n,1)
    anchor_area = anchors[..., 0] * anchors[..., 1]   # (1,9)
    iou = intersect_area / (box_area + anchor_area - intersect_area)  # (n,9)
    best_anchor = np.argmax(iou, axis=-1)  # (n,) 每个box的最大iou匹配的anchor的位置
    # 小的尺度采用大的特征图
    for t, n in enumerate(best_anchor):
        for l in range(3):
            if n in anchor_mask[l]:
                i = np.floor(box[t, 0] * grid_shapes[l]).astype('int32')
                j = np.floor(box[t, 1] * grid_shapes[l]).astype('int32')
                k = anchor_mask[l].index(n)
                c = box[t, 4].astype('int32')
                y_true[l][j, i, k, 0:4] = box[t, 0:4]
                y_true[l][j, i, k, 4] = 1
                y_true[l][j, i, k, 5 + c] = 1
    return y_true



if __name__ == '__main__':

    anchor = [[1, 1], [1, 3], [3, 2], [3, 6], [6, 4], [6, 12], [12, 9], [16, 21], [40, 35]]
    dataset = myDataSet(r'data', anchor)
    dataloader = DataLoader(dataset, batch_size=2,shuffle=True)
    for index,data in enumerate(dataloader, 0):
        print(data[0].shape)
        print(data[1].shape)
        print(data[2].shape)
        print(data[3].shape)
        break

